#!/usr/bin/python3
import os
import cv2
import numpy as np
import struct

face_attr_list_file = r'/rootfs/media/kasim/Data1/data/VideoGenderCropFace/video_crop_face_list.attr.txt'
video_label_list_file = r'/rootfs/media/kasim/Data1/data/VideoGenderCropFace/video_label_list.txt'
file_list_path = r'/rootfs/media/kasim/Data1/data/VideoGenderCropFace/BoxFeature.txt'
feature_file_path = r'/rootfs/media/kasim/Data1/data/VideoGenderCropFace/Feature.dat'

feat_dim = 256

# VERY_GOOD = 0
# GOOD = 1
# BAD = 2
# VERY_BAD = 3
FILTER_QUALITY = 3
FILTER_QUALITY_THR = 0.25
QUALITY_THR = 0.3

MIN_PITCH = -20.0
MAX_PITCH = 20.0
MIN_YAW = -20.0
MAX_YAW = 20.0
MIN_ROLL = -30.0
MAX_ROLL = 30.0

ANGLE_SCORE = 0.1
ANGLE_DIFF = 10


def cal_angle(bin_angle, k=5):
    # bin_angle = np.array(bin_angle, dtype=np.float32)
    bin_angle_index0 = np.argmax(bin_angle, 1)
    idxes = np.arange(0, bin_angle.shape[0])
    bin_angle_prob0 = bin_angle[idxes, bin_angle_index0]

    end_idx = bin_angle.shape[1] - 1

    bin_angle_probs = []
    bin_angle_indexes = []
    c = 1 + k // 2
    for offset in range(1, c):
        bin_angle_index1 = bin_angle_index0 - offset
        bin_angle_index2 = bin_angle_index0 + offset

        start_mask = bin_angle_index1 < 0
        end_mask = bin_angle_index2 > end_idx

        bin_angle_index1[start_mask] = 0
        bin_angle_index2[end_mask] = 0

        bin_angle_prob1 = bin_angle[idxes, bin_angle_index1]
        bin_angle_prob2 = bin_angle[idxes, bin_angle_index2]

        bin_angle_prob1[start_mask] = 0
        bin_angle_prob2[end_mask] = 0

        bin_angle_indexes.append(bin_angle_index1)
        bin_angle_indexes.append(bin_angle_index2)
        bin_angle_probs.append(bin_angle_prob1)
        bin_angle_probs.append(bin_angle_prob2)

    bin_angle_indexes.append(bin_angle_index0)
    bin_angle_probs.append(bin_angle_prob0)

    bin_angle_prob = np.stack(bin_angle_probs, 1)
    bin_angle_prob_norm = np.sum(bin_angle_prob, 1, keepdims=True)
    bin_angle_prob = bin_angle_prob / bin_angle_prob_norm

    bin_angle_index = np.stack(bin_angle_indexes, 1)
    bin_angle = 3 * bin_angle_index - 99

    bin_angle = bin_angle_prob * bin_angle
    # bin_angle = np.sum(bin_angle, 1, keepdims=True)
    bin_angle = np.sum(bin_angle, 1)

    angle_score = bin_angle_prob0

    return bin_angle, angle_score


def read_features(file_list, feature_file, feat_dim=256, skip_offset=8):
    file_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 1000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # gd/zs/sqqjd/ycjy/150100414a54443452064d2742d03600/video/20200509_173513/00110_000.jpg 21,38,160,224,1,1032
            lines = line.strip().split()
            file_name = lines[0].strip()

            feature = feat_file.read(feat_dim * 4 + skip_offset)
            feature = feature[skip_offset:]

            file_features[file_name] = feature

        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()
    return file_features


def main():
    face_id_images = {}
    image_list = []
    with open(face_attr_list_file, 'r') as file:
        file_count = 0
        for line in file.readlines():
            # gd/zs/nqjd/dxjyjy/150100414a54443452064d2742c93600/video/20190817_064924/00007_000.jpg
            file_path = line.strip()
            image_list.append(file_path)
            face_id = os.path.dirname(file_path)
            images = face_id_images.get(face_id, None)
            if images is None:
                face_id_images[face_id] = [file_path]
            else:
                images.append(file_path)
            file_count += 1
            # if file_count > 1000:
            #     break

    _base_path = os.path.splitext(face_attr_list_file)[0]
    title_file_path = _base_path + '.title.txt'
    with open(title_file_path, 'r') as file:
        title_str = file.readline()
        attr_item_count_str = file.readline()
        titles = title_str.strip().split(';')
        attr_item_counts = attr_item_count_str.strip().split(';')
        attr_item_counts = list(map(int, attr_item_counts))
        total_attr_item_count = sum(attr_item_counts)
        attr_count = len(titles)

    attr_item_offsets = [0]
    offset = 0
    for attr_item_count in attr_item_counts:
        offset += attr_item_count
        attr_item_offsets.append(offset)
    attrs_size = total_attr_item_count * 4
    data_file_path = _base_path + '.dat'

    image_attributes_dict = {}
    with open(data_file_path, 'rb') as file:
        for image_name in image_list:
            dat = file.read(attrs_size)
            datas = struct.unpack('{}f'.format(total_attr_item_count), dat)

            attribute = {}
            for i, title in enumerate(titles):
                start_offset = attr_item_offsets[i]
                end_offset = attr_item_offsets[i+1]
                attribute[title] = datas[start_offset:end_offset]

            image_attributes_dict[image_name] = attribute

    faceid_set = set()
    with open(video_label_list_file, 'r') as file:
        for line in file.readlines():
            video_name = line.split()[0].strip()
            faceid_set.add(video_name)

    faceid_typical_images = {}
    for face_id, images in face_id_images.items():
        if face_id not in faceid_set:
            continue
        file_names = []
        qualitys = []
        pitchs = []
        yaws = []
        rolls = []
        bin_pitchs = []
        bin_yaws = []
        bin_rolls = []
        genders = []
        for image in images:
            attributes = image_attributes_dict[image]
            quality = attributes['quality']
            pitch = attributes['pitch'][0]
            yaw = attributes['yaw'][0]
            roll = attributes['roll'][0]
            bin_pitch = attributes['bin_pitch']
            bin_yaw = attributes['bin_yaw']
            bin_roll = attributes['bin_roll']
            gender = attributes['gender']
            file_names.append(image)
            qualitys.append(quality)
            pitchs.append(pitch)
            yaws.append(yaw)
            rolls.append(roll)
            bin_pitchs.append(bin_pitch)
            bin_yaws.append(bin_yaw)
            bin_rolls.append(bin_roll)
            genders.append(gender)

        qualitys = np.array(qualitys, dtype=np.float32)
        pitchs = np.array(pitchs, dtype=np.float32)
        yaws = np.array(yaws, dtype=np.float32)
        rolls = np.array(rolls, dtype=np.float32)
        bin_pitchs = np.array(bin_pitchs, dtype=np.float32)
        bin_yaws = np.array(bin_yaws, dtype=np.float32)
        bin_rolls = np.array(bin_rolls, dtype=np.float32)
        # genders = np.array(genders, dtype=np.float32)

        # max_idx = np.argmax(qualitys, axis=1)
        # quality_mask = (max_idx == FILTER_QUALITY) & (qualitys[max_idx] >= FILTER_QUALITY_THR)

        qualitys = 0.3333333333 * qualitys[:, 2] + 0.66666666 * qualitys[:, 1] + qualitys[:, 0]
        quality_mask = qualitys >= QUALITY_THR

        bin_pitchs, bin_pitch_scores = cal_angle(bin_pitchs)
        bin_rolls, bin_roll_scores = cal_angle(bin_rolls)
        bin_yaws, bin_yaw_scores = cal_angle(bin_yaws)

        pitch_diffs = np.abs(bin_pitchs - pitchs)
        roll_diffs = np.abs(bin_rolls - rolls)
        yaw_diffs = np.abs(bin_yaws - yaws)

        score_mask = (bin_pitch_scores > ANGLE_SCORE) & (bin_roll_scores > ANGLE_SCORE) & (bin_yaw_scores > ANGLE_SCORE)
        diff_mask = (pitch_diffs < ANGLE_DIFF) & (roll_diffs < ANGLE_DIFF) & (yaw_diffs < ANGLE_DIFF)

        pitch_mask = (MIN_PITCH <= pitchs) & (pitchs <= MAX_PITCH)
        roll_mask = (MIN_ROLL <= rolls) & (rolls <= MAX_ROLL)
        yaw_mask = (MIN_YAW <= yaws) & (yaws <= MAX_YAW)
        angle_mask = pitch_mask & roll_mask & yaw_mask

        mask = score_mask & diff_mask & quality_mask & angle_mask
        if mask.any():
            # pitchs = pitchs[mask]
            # rolls = rolls[mask]
            # yaws = yaws[mask]
            # qualitys = qualitys[mask]
            # genders = genders[mask]
            file_names = [f for m, f in zip(mask, file_names) if m]
            if len(file_names) > 0:
                faceid_typical_images[face_id] = file_names

    file_features = read_features(file_list_path, feature_file_path, feat_dim=feat_dim)

    out_path = os.path.dirname(face_attr_list_file)

    out_list_file_path = os.path.join(out_path, 'Feature.select.txt')
    out_list_file = open(out_list_file_path, 'w')
    out_feat_file_path = os.path.join(out_path, 'Feature.select.dat')
    out_feat_file = open(out_feat_file_path, 'wb')
    out_file_path = os.path.join(out_path, 'face_image_list_good.txt')
    out_file = open(out_file_path, 'w')
    file_count = 0
    faceid_count = 0
    faceids = list(faceid_typical_images.keys())
    faceids.sort()
    for idx, faceid in enumerate(faceids):
        typical_images = faceid_typical_images[faceid]
        faceid_count += 1
        typical_images.sort()
        for typical_image in typical_images:
            feature = file_features.get(typical_image, None)
            if feature is None:
                continue
            file_count += 1
            out_info = '{}\n'.format(typical_image)
            out_file.write(out_info)
            feat = np.frombuffer(feature, dtype=np.float32, count=feat_dim)
            feat = np.asarray(feat) / np.linalg.norm(feat, axis=0, keepdims=True)
            feat.tofile(out_feat_file)
            out_info = '{} {} {} 0\n'.format(idx, idx, typical_image)
            out_list_file.write(out_info)
            if file_count % 1000 == 0:
                print('Proccess File Count:', file_count)
    out_file.close()
    out_feat_file.close()
    out_list_file.close()
    os.system('chmod a+wr {}'.format(out_file_path))
    os.system('chmod a+wr {}'.format(out_feat_file_path))
    os.system('chmod a+wr {}'.format(out_list_file_path))
    print('Total Face Count: {}, Filter Face Count: {}, File Count: {}'.format(len(faceid_set), faceid_count, file_count))


if __name__ == '__main__':
    main()
